{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ee90a314-aa7b-4320-ab73-241a91114a80",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "初始化设置\n",
    "\"\"\"\n",
    "import os\n",
    "import re\n",
    "import json\n",
    "import clip\n",
    "import faiss\n",
    "import torch\n",
    "import random\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from PIL import Image\n",
    "Image.MAX_IMAGE_PIXELS = None\n",
    "\n",
    "### 初始化一个clip模型\n",
    "clip_model, preprocess = clip.load(\"ViT-B/16\", device=\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7daff08d-d3a8-472a-a693-5ac429b4f765",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3334\n",
      "100\n"
     ]
    }
   ],
   "source": [
    "trainset_file_path = \"./fgvc_database/trainset.txt\"\n",
    "trainset = []\n",
    "with open(trainset_file_path, 'r') as file:\n",
    "    lines = file.readlines()\n",
    "    for line in lines:\n",
    "        split_line = line.strip().split(' ', 1)\n",
    "        part1, part2 = split_line\n",
    "        # 使用正则表达式查找数字\n",
    "        match = re.search(r'\\d+', part2)\n",
    "        label = int(match.group())\n",
    "        trainset.append([part1, label])\n",
    "print(len(trainset))\n",
    "\n",
    "classnames_file_path = \"./fgvc_database/classnames.txt\"\n",
    "with open(classnames_file_path, 'r') as file:\n",
    "    classnames = file.readlines()\n",
    "print(len(classnames))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "f1756807-533b-46b7-b27e-ae03dd4e54fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "使用random的方法产生的10000条数据，Yes5000条，No5000条\n",
    "\"\"\"\n",
    "json_data = []\n",
    "count = 0 \n",
    "while count<5000:\n",
    "    sampled_elements = random.sample(trainset, 2)\n",
    "    sampled_name = [classnames[int(sampled_elements[0][1])].strip(), classnames[int(sampled_elements[1][1])].strip()]\n",
    "    if sampled_elements[0][1] != sampled_elements[1][1]:\n",
    "        my_dict = {\n",
    "        'id': i,\n",
    "        'image': [sampled_elements[0][0], sampled_elements[1][0]],\n",
    "        'conversations': [{'from': 'user', 'value': '以下两张图片中的飞机，是否属于同一型号<Img index=0><image></Img> <Img index=1><image></Img>请直接给出Yes或No的回答，再分析原因'},{'from': 'assistant', 'value': 'No. <InsertImg index=0>是xxx1型号，而 <InsertImg index=1>是xxx2型号'}]\n",
    "        }\n",
    "        my_dict[\"conversations\"][1][\"value\"] = my_dict[\"conversations\"][1][\"value\"].replace('xxx1', sampled_name[0])\n",
    "        my_dict[\"conversations\"][1][\"value\"] = my_dict[\"conversations\"][1][\"value\"].replace('xxx2', sampled_name[1])\n",
    "        json_data.append(my_dict)\n",
    "        count+=1\n",
    "count = 0\n",
    "while count<5000:\n",
    "    sampled_elements = random.sample(trainset, 2)\n",
    "    sampled_name = [classnames[int(sampled_elements[0][1])].strip(), classnames[int(sampled_elements[1][1])].strip()]\n",
    "    if sampled_elements[0][1] == sampled_elements[1][1]:\n",
    "        my_dict = {\n",
    "        'id': i,\n",
    "        'image': [sampled_elements[0][0], sampled_elements[1][0]],\n",
    "        'conversations': [{'from': 'user', 'value': '以下两张图片中的飞机，是否属于同一型号<Img index=0><image></Img> <Img index=1><image></Img>请直接给出Yes或No的回答，再分析原因'},{'from': 'assistant', 'value': 'Yes. <InsertImg index=0>和<InsertImg index=1>都是xxx型号'}]\n",
    "        }\n",
    "        my_dict[\"conversations\"][1][\"value\"] = my_dict[\"conversations\"][1][\"value\"].replace('xxx', sampled_name[0])\n",
    "        json_data.append(my_dict)\n",
    "        count+=1\n",
    "        \n",
    "with open('./my_list.json', 'w', encoding='utf-8') as file:\n",
    "    json.dump(json_data, file, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "90e6e371-ced2-45b6-9e94-798a830bdadc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3334/3334 [01:25<00:00, 39.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "39282\n",
      "133360\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "使用clip的方法产生的10000条数据\n",
    "\"\"\"\n",
    "### 初始化\n",
    "index_img_save_path = \"./fgvc_database/index_img.index\"\n",
    "index = faiss.read_index(index_img_save_path)\n",
    "\n",
    "json_data = []\n",
    "yes_number = 0\n",
    "no_number = 0\n",
    "### 遍历数据集\n",
    "for train_data in tqdm(trainset):\n",
    "    \n",
    "    ### 获取每一个数据的图片位置和label\n",
    "    train_data_label = int(train_data[1])\n",
    "    train_data_image = train_data[0]\n",
    "    # print(train_data_label)\n",
    "    # print(train_data_image)\n",
    "\n",
    "    ### 匹配临近的图片\n",
    "    with torch.no_grad():\n",
    "        image = preprocess(Image.open(train_data[0])).unsqueeze(0).to(\"cuda\")\n",
    "        # torch.Size([1, 512])\n",
    "        image_features = clip_model.encode_image(image)\n",
    "        image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "        image_features = np.array(image_features.cpu())\n",
    "        distance, index_result = index.search(image_features, 100)\n",
    "        # print(\"Distance:\")\n",
    "        # print(distance)\n",
    "        # print(index_result[0])\n",
    "    \n",
    "        matched_labels = []\n",
    "        matched_images = []\n",
    "        with open(trainset_file_path, 'r') as file:\n",
    "            lines = file.readlines()\n",
    "            for index_number in index_result[0]:\n",
    "                parts = lines[index_number].strip().split(' ', 1)\n",
    "                # print(parts)\n",
    "                part1, part2 = parts\n",
    "                match = re.search(r\":\\s*(\\d+)\", part2)\n",
    "                number = int(match.group(1))\n",
    "                matched_labels.append(number)\n",
    "                matched_images.append(part1)\n",
    "        # print(matched_labels) \n",
    "        # print(matched_images)\n",
    "\n",
    "    ### 遍历匹配的图片。构造数据\n",
    "    count = 0 \n",
    "    for i in range(99):\n",
    "        if count<40:\n",
    "            if train_data_label != matched_labels[i+1]:\n",
    "                train_data_name = classnames[train_data_label].strip()\n",
    "                matched_data_name = classnames[matched_labels[i+1]].strip()\n",
    "                temp_train_data_image = train_data_image.replace(\"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/CLIP-Cls/data/\",\"/mnt/petrelfs/share_data/zangyuhang/img-cls/\")\n",
    "                temp_matched_images = matched_images[i+1].replace(\"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/CLIP-Cls/data/\",\"/mnt/petrelfs/share_data/zangyuhang/img-cls/\")\n",
    "                my_dict = {\n",
    "                'id': int(len(json_data)),\n",
    "                'image': [temp_train_data_image, temp_matched_images],\n",
    "                'conversations': [{'from': 'user', 'value': '以下两张图片中的飞机，是否属于同一型号<Img index=0><image></Img> <Img index=1><image></Img>请直接给出Yes或No的回答，再分析原因'},{'from': 'assistant', 'value': 'No. <InsertImg index=0>是xxx1型号，而 <InsertImg index=1>是xxx2型号'}]\n",
    "                }\n",
    "                my_dict[\"conversations\"][1][\"value\"] = my_dict[\"conversations\"][1][\"value\"].replace('xxx1', train_data_name)\n",
    "                my_dict[\"conversations\"][1][\"value\"] = my_dict[\"conversations\"][1][\"value\"].replace('xxx2', matched_data_name)\n",
    "                json_data.append(my_dict)\n",
    "                count+=1\n",
    "                no_number+=1\n",
    "\n",
    "    count = 0 \n",
    "    for i in range(99):\n",
    "        if count<40:\n",
    "            if train_data_label == matched_labels[i+1]:\n",
    "                train_data_name = classnames[train_data_label].strip()\n",
    "                matched_data_name = classnames[matched_labels[i+1]].strip()\n",
    "                temp_train_data_image = train_data_image.replace(\"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/CLIP-Cls/data/\",\"/mnt/petrelfs/share_data/zangyuhang/img-cls/\")\n",
    "                temp_matched_images = matched_images[i+1].replace(\"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/CLIP-Cls/data/\",\"/mnt/petrelfs/share_data/zangyuhang/img-cls/\")\n",
    "                my_dict = {\n",
    "                'id': int(len(json_data)),\n",
    "                'image': [temp_train_data_image, temp_matched_images],\n",
    "                'conversations': [{'from': 'user', 'value': '以下两张图片中的飞机，是否属于同一型号<Img index=0><image></Img> <Img index=1><image></Img>请直接给出Yes或No的回答，再分析原因'},{'from': 'assistant', 'value': 'Yes. <InsertImg index=0>和<InsertImg index=1>都是xxx型号'}]\n",
    "                }\n",
    "                my_dict[\"conversations\"][1][\"value\"] = my_dict[\"conversations\"][1][\"value\"].replace('xxx', train_data_name)\n",
    "                json_data.append(my_dict)\n",
    "                count+=1\n",
    "                yes_number+=1\n",
    "    # print(json_data)\n",
    "print(yes_number)\n",
    "print(no_number)\n",
    "\n",
    "with open('./clip_genetator.json', 'w', encoding='utf-8') as file:\n",
    "    json.dump(json_data, file, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "28f1d105-36c3-4b9f-b361-11d9216bc81a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "130575\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "确保图像对独一无二，防止出现相同的数据\n",
    "\"\"\"\n",
    "with open('./clip_genetator.json', 'r', encoding='utf-8') as file:\n",
    "    data = json.load(file)\n",
    "\n",
    "unique_image_pairs = set()\n",
    "new_data = []\n",
    "\n",
    "# 检查并移除非唯一图像对的元素\n",
    "for item in data:\n",
    "    image_pair = tuple(sorted(item['image']))\n",
    "    if image_pair not in unique_image_pairs:\n",
    "        unique_image_pairs.add(image_pair)\n",
    "        new_data.append(item)\n",
    "\n",
    "# 重新命名id\n",
    "for i in range(len(new_data)):\n",
    "    new_data[i][\"id\"] = i\n",
    "\n",
    "print(len(new_data))\n",
    "\n",
    "# 将处理后的数据写入新文件\n",
    "with open(\"cleaned.json\", 'w', encoding='utf-8') as file:\n",
    "    json.dump(new_data, file, indent=4, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b4cebcd3-522d-413f-b271-943a8619f931",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3334/3334 [01:21<00:00, 40.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "30647\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "使用clip的方法,生成图文排序的数据集\n",
    "\"\"\"\n",
    "def custom_sort(item, match):\n",
    "    return 0 if item == match else 1\n",
    "    \n",
    "### 初始化\n",
    "index_img_save_path = \"./fgvc_database/index_img.index\"\n",
    "index = faiss.read_index(index_img_save_path)\n",
    "\n",
    "json_data = []\n",
    "### 遍历数据集\n",
    "for train_data in tqdm(trainset):\n",
    "    \n",
    "    ### 获取每一个数据的图片位置和label\n",
    "    train_data_label = int(train_data[1])\n",
    "    train_data_image = train_data[0]\n",
    "    train_data_name = classnames[train_data_label].strip()\n",
    "    # print(train_data_label)\n",
    "    # print(train_data_image)\n",
    "    # print(train_data_name)\n",
    "\n",
    "    ### 匹配临近的图片\n",
    "    with torch.no_grad():\n",
    "        image = preprocess(Image.open(train_data[0])).unsqueeze(0).to(\"cuda\")\n",
    "        # torch.Size([1, 512])\n",
    "        image_features = clip_model.encode_image(image)\n",
    "        image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "        image_features = np.array(image_features.cpu())\n",
    "        distance, index_result = index.search(image_features, 20)\n",
    "        # print(\"Distance:\")\n",
    "        # print(distance)\n",
    "        # print(index_result[0])\n",
    "    \n",
    "        matched_labels = []\n",
    "        matched_images = []\n",
    "        matched_names = []\n",
    "        with open(trainset_file_path, 'r') as file:\n",
    "            lines = file.readlines()\n",
    "            for index_number in index_result[0]:\n",
    "                parts = lines[index_number].strip().split(' ', 1)\n",
    "                # print(parts)\n",
    "                part1, part2 = parts\n",
    "                match = re.search(r\":\\s*(\\d+)\", part2)\n",
    "                number = int(match.group(1))\n",
    "                matched_labels.append(number)\n",
    "                matched_images.append(part1)\n",
    "                matched_names.append(classnames[number].strip())\n",
    "        # print(matched_labels) \n",
    "        # print(matched_images)\n",
    "        # print(matched_names)\n",
    "\n",
    "        ### 五个五个拆分,且确保每一个lsit中存在正确答案\n",
    "        grouped_lists = []\n",
    "        grouped_list = [matched_names[i:i + 5] for i in range(0, len(matched_names) - 4)]\n",
    "        for i in range(len(grouped_list)):\n",
    "            if train_data_name in grouped_list[i]:\n",
    "                grouped_lists.append(grouped_list[i])\n",
    "        # print(grouped_lists)\n",
    "\n",
    "        ### 重新排序，生成正确答案\n",
    "        sorted_lists = []\n",
    "        for i in range(len(grouped_lists)):\n",
    "            sorted_list = sorted(grouped_lists[i], key=lambda item: custom_sort(item, train_data_name))\n",
    "            sorted_lists.append(sorted_list)\n",
    "        # print(sorted_lists)\n",
    "\n",
    "        ### 转化为string\n",
    "        grouped_lists = [str(sublist) for sublist in grouped_lists]\n",
    "        sorted_lists = [str(sublist) for sublist in sorted_lists]\n",
    "        \n",
    "        for i in range(len(grouped_lists)):\n",
    "            my_dict = {\n",
    "                    'id': int(len(json_data)),\n",
    "                    'image': [train_data_image],\n",
    "                    'conversations': [{'from': 'user', 'value': 'Here is a image:<Img index=1><image></Img>. Please play the role of a aircraft classification expert, and sort the provided categories from high to low according to the top 5 similarity with the input image. Here are the optional categories:{names}.'.format(names=grouped_lists[i])},{'from': 'assistant', 'value': '{result}'.format(result=sorted_lists[i])}]\n",
    "                    }\n",
    "            json_data.append(my_dict)\n",
    "\n",
    "        # print(json_data)\n",
    "print(len(json_data))\n",
    "with open('./classnames_rerank.json', 'w', encoding='utf-8') as file:\n",
    "    json.dump(json_data, file, indent=4, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "526e9470-59b0-4f8a-a46d-6c8cd37f94ad",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
